import numpy as np
from local_search_operators import get_operator
from problems.TSP.fitness import calculate_fitness


def tsp_apply_local_search(new_population, problem_instance, elite_size, ls_intensity, operator_names):
    """
    TSP特定的局部搜索应用逻辑

    参数:
    new_population: 当前种群
    problem_instance: 问题实例
    elite_size: 精英个体数量
    ls_intensity: 局部搜索强度
    operator_names: 要使用的算子名称列表

    返回:
    应用局部搜索后的种群
    """
    fitness_values = calculate_fitness(new_population, problem_instance.cities)

    elite_indices = np.argsort(fitness_values)[-elite_size:]

    if operator_names is None:
        operator_names = ["two_opt"]

    # 对每个精英个体应用局部搜索
    for idx in elite_indices:
        current_solution = new_population[idx]

        # 依次应用所有指定的算子
        for op_name in operator_names:
            operator = get_operator(op_name)
            if operator:
                improved_solution, improvement = operator(
                    current_solution, problem_instance, ls_intensity
                )
                current_solution = improved_solution

        new_population[idx] = current_solution

    return new_population


